{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8639ad23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# type: ignore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b237311e",
   "metadata": {},
   "source": [
    "# Together Supervised Fine-Tuning\n",
    "\n",
    "This recipe allows TensorZero users to fine-tune Together models using their own data.\n",
    "Since TensorZero automatically logs all inferences and feedback, it is straightforward to fine-tune a model using your own data and any prompt you want.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18b7f8ce",
   "metadata": {},
   "source": [
    "To get started:\n",
    "\n",
    "- Set the `TENSORZERO_CLICKHOUSE_URL` environment variable. For example: `TENSORZERO_CLICKHOUSE_URL=\"http://chuser:chpassword@localhost:8123/tensorzero\"`\n",
    "- Set the `TOGETHER_API_KEY` environment variable.\n",
    "- Update the following parameters:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "563de2b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONFIG_PATH = \"../../../examples/data-extraction-ner/config/tensorzero.toml\"\n",
    "\n",
    "FUNCTION_NAME = \"extract_entities\"\n",
    "\n",
    "METRIC_NAME = \"jaccard_similarity\"\n",
    "\n",
    "# The name of the variant to use to grab the templates used for fine-tuning\n",
    "TEMPLATE_VARIANT_NAME = \"gpt_4o_mini\"  # It's OK that this variant uses a different model than the one we're fine-tuning\n",
    "\n",
    "# If the metric is a float metric, you can set the threshold to filter the data\n",
    "FLOAT_METRIC_THRESHOLD = 0.5\n",
    "\n",
    "# Fraction of the data to use for validation\n",
    "VAL_FRACTION = 0.2\n",
    "\n",
    "# Maximum number of samples to use for fine-tuning\n",
    "MAX_SAMPLES = 100_000\n",
    "\n",
    "# The name of the model to fine-tune (supported models: https://docs.together.ai/docs/fine-tuning-models)\n",
    "MODEL_NAME = \"meta-llama/Meta-Llama-3.1-8B-Instruct-Reference\"\n",
    "\n",
    "# At the time of writing, Together.ai does not support tool call content blocks in assistant messages. Or the tool role.\n",
    "# We will drop these invalid messages from the dataset by default.\n",
    "# You can set this to False to keep the invalid messages in the dataset.\n",
    "DROP_INVALID_MESSAGES = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec7ccefd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "tensorzero_path = os.path.abspath(os.path.join(os.getcwd(), \"../../../\"))\n",
    "if tensorzero_path not in sys.path:\n",
    "    sys.path.append(tensorzero_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6730012c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import subprocess\n",
    "import tempfile\n",
    "import time\n",
    "from pprint import pprint\n",
    "from typing import Any, Dict, List\n",
    "\n",
    "import requests\n",
    "import toml\n",
    "from IPython.display import clear_output\n",
    "from tensorzero import (\n",
    "    FloatMetricFilter,\n",
    "    TensorZeroGateway,\n",
    ")\n",
    "\n",
    "from recipes.util import tensorzero_rendered_samples_to_conversations, train_val_split"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e76a20fc",
   "metadata": {},
   "source": [
    "Initialize the TensorZero client\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb3320ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorzero_client = TensorZeroGateway.build_embedded(\n",
    "    config_file=CONFIG_PATH,\n",
    "    clickhouse_url=os.environ[\"TENSORZERO_CLICKHOUSE_URL\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9952cf24",
   "metadata": {},
   "source": [
    "Set the metric filter as needed\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "084e8e80",
   "metadata": {},
   "outputs": [],
   "source": [
    "comparison_operator = \">=\"\n",
    "metric_node = FloatMetricFilter(\n",
    "    metric_name=METRIC_NAME,\n",
    "    value=FLOAT_METRIC_THRESHOLD,\n",
    "    comparison_operator=comparison_operator,\n",
    ")\n",
    "# from tensorzero import BooleanMetricFilter\n",
    "# metric_node = BooleanMetricFilter(\n",
    "#     metric_name=METRIC_NAME,\n",
    "#     value=True  # or False\n",
    "# )\n",
    "\n",
    "metric_node"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e2b9ce5",
   "metadata": {},
   "source": [
    "Query the inferences from ClickHouse\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "714ddecb",
   "metadata": {},
   "outputs": [],
   "source": [
    "stored_inferences = tensorzero_client.experimental_list_inferences(\n",
    "    function_name=FUNCTION_NAME,\n",
    "    variant_name=None,\n",
    "    output_source=\"inference\",  # could also be \"demonstration\"\n",
    "    filters=metric_node,\n",
    "    limit=MAX_SAMPLES,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56fe3524",
   "metadata": {},
   "source": [
    "Render the inputs using the templates in the template variant.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b4a45ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "rendered_samples = tensorzero_client.experimental_render_samples(\n",
    "    stored_inferences=stored_inferences,\n",
    "    variants={FUNCTION_NAME: TEMPLATE_VARIANT_NAME},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a9c7bf4",
   "metadata": {},
   "source": [
    "Split the data into training and validation sets for fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69a01b38",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_samples, val_samples = train_val_split(\n",
    "    rendered_samples,\n",
    "    val_size=VAL_FRACTION,\n",
    "    last_inference_only=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2f6245c",
   "metadata": {},
   "source": [
    "Convert the rendered samples to openai format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64cfd501",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_samples = tensorzero_rendered_samples_to_conversations(\n",
    "    train_samples, conversation_key=\"messages\", join_text_blocks=True\n",
    ")\n",
    "val_samples = tensorzero_rendered_samples_to_conversations(\n",
    "    val_samples, conversation_key=\"messages\", join_text_blocks=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "718b4033",
   "metadata": {},
   "source": [
    "We'll write the training and validation messages to temporary files for the Together CLI\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1c4b5c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def upload_dataset_to_together(samples: List[Dict[str, Any]]) -> str:\n",
    "    with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".jsonl\", delete=False) as f:\n",
    "        # Write the conversational_messages to the temporary file\n",
    "        for item in samples:\n",
    "            json.dump(item, f)\n",
    "            f.write(\"\\n\")\n",
    "        f.flush()\n",
    "\n",
    "        dataset_path = f.name\n",
    "        result = subprocess.run([\"together\", \"files\", \"upload\", dataset_path], capture_output=True)\n",
    "        print(\"Stdout:\")\n",
    "        print(result.stdout.decode())\n",
    "        print(\"Stderr:\")\n",
    "        print(result.stderr.decode())\n",
    "        together_result = json.loads(result.stdout)\n",
    "        return together_result[\"id\"]\n",
    "\n",
    "\n",
    "train_file_object_id = upload_dataset_to_together(train_samples)\n",
    "val_file_object_id = upload_dataset_to_together(val_samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d91ba000",
   "metadata": {},
   "source": [
    "Launch the fine-tuning job.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "519197b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://api.together.xyz/v1/fine-tunes\"\n",
    "print(\"MODEL: \", MODEL_NAME)\n",
    "print(\"Train: \", train_file_object_id)\n",
    "print(\"Val: \", val_file_object_id)\n",
    "\n",
    "payload = {\n",
    "    \"training_file\": train_file_object_id,\n",
    "    \"validation_file\": val_file_object_id,\n",
    "    \"model\": MODEL_NAME,\n",
    "    \"n_epochs\": 1,\n",
    "    \"n_checkpoints\": 1,\n",
    "    \"n_evals\": 0,\n",
    "    \"batch_size\": 16,\n",
    "    \"learning_rate\": 0.00001,\n",
    "    \"lr_scheduler\": {\"lr_scheduler_args\": {\"min_lr_ratio\": 0}},\n",
    "    \"warmup_ratio\": 0,\n",
    "    \"max_grad_norm\": 1,\n",
    "    \"weight_decay\": 0,\n",
    "    \"train_on_inputs\": \"auto\",\n",
    "    \"training_type\": {\"type\": \"Lora\", \"lora_r\": 8, \"lora_alpha\": 32},\n",
    "}\n",
    "headers = {\n",
    "    \"accept\": \"application/json\",\n",
    "    \"content-type\": \"application/json\",\n",
    "    \"authorization\": f\"Bearer {os.environ['TOGETHER_API_KEY']}\",\n",
    "}\n",
    "\n",
    "response = requests.post(url, json=payload, headers=headers)\n",
    "print(\"Response status: \", response.status_code)\n",
    "print(\"Response body: \")\n",
    "print(response.text)\n",
    "response_json = json.loads(response.text)\n",
    "fine_tune_id = response_json[\"id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d47e7f1",
   "metadata": {},
   "source": [
    "Wait for the fine-tuning job to complete.\n",
    "\n",
    "This cell will take a while to run.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7876d76",
   "metadata": {},
   "outputs": [],
   "source": [
    "while True:\n",
    "    clear_output(wait=True)\n",
    "\n",
    "    try:\n",
    "        job_status = requests.get(\n",
    "            f\"https://api.together.xyz/v1/fine-tunes/{fine_tune_id}\",\n",
    "            headers={\n",
    "                \"accept\": \"application/json\",\n",
    "                \"authorization\": f\"Bearer {os.environ['TOGETHER_API_KEY']}\",\n",
    "            },\n",
    "        ).json()\n",
    "        pprint(job_status)\n",
    "        print(\"Status: \", job_status[\"status\"])\n",
    "        if job_status[\"status\"] in (\"completed\", \"failed\", \"cancelled\"):\n",
    "            break\n",
    "    except Exception as e:\n",
    "        print(f\"Error: {e}\")\n",
    "\n",
    "    time.sleep(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da24ba74",
   "metadata": {},
   "source": [
    "Once the fine-tuning job is complete, you can add the fine-tuned model to your config file.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ec6fdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "fine_tuned_model = job_status[\"model_output_name\"]\n",
    "model_config = {\n",
    "    \"models\": {\n",
    "        fine_tuned_model: {\n",
    "            \"routing\": [\"together\"],\n",
    "            \"providers\": {\"together\": {\"type\": \"together\", \"model_name\": fine_tuned_model}},\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "print(toml.dumps(model_config))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45ba392c",
   "metadata": {},
   "source": [
    "Finally, add a new variant to your function to use the fine-tuned model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "136faf8d",
   "metadata": {},
   "source": [
    "You're all set!\n",
    "\n",
    "You can change the weight to enable a gradual rollout of the new model.\n",
    "\n",
    "You might also add other parameters (e.g. `max_tokens`, `temperature`) to the variant section in the config file.\n"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "formats": "ipynb,py:percent",
   "main_language": "python"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
